package pg

import (
	"context"
	"io"
	"os"
	"path/filepath"
	"slices"
	"strings"
	"testing"

	"github.com/stretchr/testify/require"
	"google.golang.org/protobuf/proto"
	"gopkg.in/yaml.v3"

	storepb "github.com/bytebase/bytebase/backend/generated-go/store"
	"github.com/bytebase/bytebase/backend/plugin/advisor"
	"github.com/bytebase/bytebase/backend/plugin/parser/base"

	// Import pg parser to register PostgreSQL parser via init()
	_ "github.com/bytebase/bytebase/backend/plugin/parser/pg"
	"github.com/bytebase/bytebase/backend/plugin/schema"
	"github.com/bytebase/bytebase/backend/store/model"
)

// TestCase is the data struct for test.
type TestCase struct {
	Statement string            `yaml:"statement"`
	EnableSDL bool              `yaml:"enableSDL"`
	Want      []*storepb.Advice `yaml:"want,omitempty"`
}

// RunPGSQLReviewRuleTest helps to test PostgreSQL SQL review rules.
// This is similar to advisor.RunSQLReviewRuleTest.
func RunPGSQLReviewRuleTest(t *testing.T, rule storepb.SQLReviewRule_Type, dbType storepb.Engine, record bool) {
	var tests []TestCase

	fileName := strings.Map(func(r rune) rune {
		switch r {
		case '.', '-':
			return '_'
		default:
			return r
		}
	}, strings.ToLower(rule.String()))
	filepath := filepath.Join("test", fileName+".yaml")
	yamlFile, err := os.Open(filepath)
	require.NoError(t, err)
	defer yamlFile.Close()

	byteValue, err := io.ReadAll(yamlFile)
	require.NoError(t, err)
	err = yaml.Unmarshal(byteValue, &tests)
	require.NoError(t, err, rule)

	for i, tc := range tests {
		// Add engine types here for mocked database metadata.
		var schemaMetadata *storepb.DatabaseSchemaMetadata
		switch dbType {
		case storepb.Engine_POSTGRES:
			schemaMetadata = advisor.MockPostgreSQLDatabase
		default:
			t.Fatalf("%s doesn't have mocked metadata support", storepb.Engine_name[int32(dbType)])
		}

		// Create OriginalMetadata as DatabaseMetadata (read-only)
		// Clone to avoid mutations affecting future test cases
		originalCatalogClone, ok := proto.Clone(schemaMetadata).(*storepb.DatabaseSchemaMetadata)
		require.True(t, ok, "failed to clone schema metadata")
		originalSchema := model.NewDatabaseMetadata(originalCatalogClone, nil, nil, dbType, true /* isCaseSensitive for PostgreSQL */)
		originalMetadata := originalSchema

		// Create FinalMetadata as DatabaseMetadata (mutable for walk-through)
		// Clone to avoid mutations affecting future test cases
		finalCatalogClone, ok := proto.Clone(schemaMetadata).(*storepb.DatabaseSchemaMetadata)
		require.True(t, ok, "failed to clone schema metadata")
		finalMetadata := model.NewDatabaseMetadata(finalCatalogClone, nil, nil, storepb.Engine_POSTGRES, true /* isCaseSensitive for PostgreSQL */)

		// Get default payload, or use empty string for test-only rules
		payload, err := advisor.SetDefaultSQLReviewRulePayload(rule, dbType)
		if err != nil {
			// For test-only rules (like "hello-world"), use empty payload
			payload = ""
		}

		// Parse SQL using base.Parse() to get AST
		unifiedASTs, err := base.Parse(dbType, tc.Statement)
		require.NoError(t, err, "Failed to parse SQL: %s", tc.Statement)

		// Skip test cases with only comments or no SQL statements (parser returns nil/empty)
		if len(unifiedASTs) == 0 {
			// For empty/comment-only statements, the advisor should return no advice
			if len(tc.Want) > 0 {
				t.Fatalf("Test case expects advice but statement has no SQL: %s", tc.Statement)
			}
			continue
		}

		// Walk through the catalog to build metadata using AST directly
		advice := schema.WalkThrough(dbType, finalMetadata, unifiedASTs)
		require.Nil(t, advice, "Failed to walk through final catalog: %s", tc.Statement)

		ruleList := []*storepb.SQLReviewRule{
			{
				Type:    rule,
				Level:   storepb.SQLReviewRule_WARNING,
				Payload: payload,
			},
		}

		// Call the advisor directly with ANTLR AST
		adviceList, err := advisor.Check(
			context.Background(),
			dbType,
			rule,
			advisor.Context{
				DBSchema:                 schemaMetadata,
				EnableSDL:                tc.EnableSDL,
				EnablePriorBackup:        true,
				AST:                      unifiedASTs, // Pass AST list
				Statements:               tc.Statement,
				Rule:                     ruleList[0],
				OriginalMetadata:         originalMetadata,
				FinalMetadata:            finalMetadata,
				Driver:                   nil,
				CurrentDatabase:          "TEST_DB",
				UsePostgresDatabaseOwner: true,
			},
		)

		// Sort adviceList by (line, content) for consistent comparison
		slices.SortFunc(adviceList, func(x, y *storepb.Advice) int {
			if x.GetStartPosition() == nil || y.GetStartPosition() == nil {
				if x.GetStartPosition() == nil && y.GetStartPosition() == nil {
					return 0
				} else if x.GetStartPosition() == nil {
					return -1
				}
				return 1
			}
			if x.GetStartPosition().Line != y.GetStartPosition().Line {
				if x.GetStartPosition().Line < y.GetStartPosition().Line {
					return -1
				}
				return 1
			}
			if x.Content < y.Content {
				return -1
			} else if x.Content > y.Content {
				return 1
			}
			return 0
		})

		require.NoError(t, err)
		if record {
			tests[i].Want = adviceList
		} else {
			require.Equalf(t, tc.Want, adviceList, "rule: %s, statements: %s", rule, tc.Statement)
		}
	}

	if record {
		err := yamlFile.Close()
		require.NoError(t, err)
		byteValue, err := yaml.Marshal(tests)
		require.NoError(t, err)
		err = os.WriteFile(filepath, byteValue, 0644)
		require.NoError(t, err)
	}
}
